# -*- encoding: utf-8 -*-

import sys
import json
import re
import time

from aibox.data import base_es_io
from elasticsearch import Elasticsearch

#连接ES
from pandas.io.json import json_normalize

es = Elasticsearch(['10.2.111.51'],http_auth=('admin','admin'),port=9200)


def se_search( startTime, endTime):

    query_json = {
        'query':{
            'bool':{
                'must':[
                    {'range':{'gcsj':{
                        "gte":"1526000400000",
                        "lte":"1526002200000"
                        # 1526000460000
                    }}}
            ]
            }
        }
    }

    all_Doc = es.search(index='ehlindex201805',doc_type='pass_car',size= 1000000 ,body=query_json)
    df = parse_data_from_json(all_Doc['hits']['hits'])
    # t = BaseESIO()
    # df = t.query(query_json=query_json)
    print(len(df))
    # for hit in all_Doc['hits']['hits']:
    #     print(hit)
        # print(hit['_source'])
def parse_data_from_json( res, field='_source.'):

    df = json_normalize(res)
    for col in df.columns:
        if field not in col:
            df.drop(col, axis=1, inplace=True)

    df.rename(columns=lambda x: x.replace(field, ''), inplace=True)
    # df = df.set_index("_id")

    return df

from aibox.api import get_current_dt
from aibox.interface import AbstractDBIO
from elasticsearch import Elasticsearch, helpers, TransportError
from pandas.io.json import json_normalize


# ES数据IO类
from aibox.utils.logger import user_system_log


class BaseESIO(AbstractDBIO):

    def __init__(self):
        self._engine = self.get_engine()
        self._dt_col = None
        self.index_func = None

    def get_engine(self, config=None):


        return es

    def query(self, startTime=None , endTime=None ,query_json=None):
        if query_json is None:
            start = startTime.tz_localize('Asia/Shanghai')
            end = endTime.tz_localize('Asia/Shanghai')

            query_json = {
                'query': {
                    'bool': {
                        'must': [
                            {'range': {self._dt_col: {  # 过车时间
                                'gte': int(start.value / 10 ** 6),
                                'lte': int(end.value / 10 ** 6),
                            }
                            }}
                        ],
                    }
                }
            }
        # es索引名称创建规则
        engine = self.get_engine()
        try:
            res = list(helpers.scan(client=engine, scroll='2m', index='ehlindex201805', doc_type='pass_car',
                                query=query_json))
        except TransportError:
            user_system_log.error('索引异常，请检查ES配置')
            return None
        except Exception as e:
            user_system_log.error(e)
            return None
        # 如果没有检索出数据，就返回None
        if len(res) == 0:
            return None
        df = self.parse_data_from_json(res)
        return df

    def create_engine(self, config):
        es = Elasticsearch(['10.2.111.51'], http_auth=('admin', config.pwd),
                           port=config.port)
        return es

    def parse_data_from_json(self, res, field='_source.'):

        df = json_normalize(res)
        for col in df.columns:
            if field not in col:
                df.drop(col, axis=1, inplace=True)

        df.rename(columns=lambda x: x.replace(field, ''), inplace=True)
        # df = df.set_index("_id")

        return df



startTime = 1496916033000
endTime = 1496917210000
start = time.clock()

se_search( startTime, endTime)
end = time.clock()
print('Running time: %s Seconds' % (end - start))